import math
import matplotlib.pyplot as plt
import torch
from params.Attention_CNN_FC_params import *
from utils.timeloader import getLoader
device = 'cpu'
from torch import nn


class MHAttention(nn.Module,):
    def __init__(self,input_dim,num_hiddens,num_head,dropout = 0.3,bias = False):
        super(MHAttention, self).__init__()
        self.W_q = nn.Linear(input_dim,num_hiddens)
        self.W_k = nn.Linear(input_dim,num_hiddens)
        self.W_v = nn.Linear(input_dim,num_hiddens)
        self.attn = nn.MultiheadAttention(embed_dim= num_hiddens,num_heads=num_head,dropout=dropout,batch_first=True)
        self.W_o = nn.Linear(num_hiddens,input_dim)

    def forward(self,x):
        attn_output,_= self.attn(self.W_q(x),self.W_k(x),self.W_v(x))
        return self.W_o(attn_output)


class AttentionBlock(nn.Module):
    def __init__(self,input_dim,num_hiddens,num_head,dropout = 0.3):
        super(AttentionBlock, self).__init__()
        self.attn = MHAttention(input_dim,num_hiddens,num_head,dropout)
        self.norm1 = nn.LayerNorm(input_dim)
        self.norm2 = nn.LayerNorm(input_dim)
        self.fc = nn.Sequential(
            nn.Linear(input_dim,input_dim * 2),
            nn.ReLU(),
            nn.Linear(input_dim * 2, input_dim)
        )

    def forward(self, x):
        x_ = self.attn(x)
        x_ = self.norm1(x + x_)
        x = self.fc(x_)
        return self.norm2(x + x_)


class SENET(nn.Module):
    def __init__(self,channel,reduction = 4):
        super(SENET, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool1d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(self,x):
        b,t,_ = x.size()
        y = self.avg_pool(x).view(b,t)
        y = self.fc(y).view(b,t,1)
        return x * y.expand_as(x)

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.cnn = nn.Sequential(
            nn.Conv1d(in_channels=1,out_channels=20,stride=4,kernel_size=17,padding=8),
            SENET(channel=20),
            nn.MaxPool1d(kernel_size=5,stride=2,padding=2),
            nn.ReLU(),
            nn.Conv1d(in_channels=20,out_channels=40,stride=2,kernel_size=5,padding=2),
            SENET(40),
            nn.MaxPool1d(kernel_size=5, stride=2, padding=2),
            nn.ReLU(),
            nn.Conv1d(in_channels=40, out_channels=1, stride=1, kernel_size=3, padding=1),
        )

    def forward(self,x):
        b = x.shape[0]
        c = x.shape[1]
        x = x.view(b * c, -1)
        x = x.unsqueeze(1)
        x = self.cnn(x)
        x = x.view(b,c,-1)
        return x


class MODEL(nn.Module):
    def __init__(self,input_size):
        super(MODEL, self).__init__()
        self.cnn = CNN()
        self.attn = nn.Sequential(
            AttentionBlock(64, 128,2),
            AttentionBlock(64, 128,2),
        )
        self.fc = nn.Linear(64,1)
        self.out = nn.Linear(30,1)
    def forward(self,x):
        x = self.cnn(x)
        x = self.attn(x)
        x = self.fc(x).squeeze(2)
        return self.out(x)




test_bearing = 7
temp_num = math.log(0.5)
model = torch.load('./74.pt').to(device)
for test_bearing in range(3,8):
    test_set = getLoader('./test_files', batch_size=batch_size, bearing_condition=bearing_condition,
                         bearing_label=test_bearing, window=window)
    pltx = []  # 实际值
    plty = []  # 预测值
    sum_RMSE_loss = 0
    sum_MAE_loss = 0
    sum_RUL_score = 0
    m = 0
    for i, (x, y, label) in enumerate(test_set):
        x = x.to(device)
        y = y.to(device)
        label_hat = model(x)  # 此时label_hat是一个二维的数组（batch_size,1)

        label_hat = label_hat.tolist()
        label = label.tolist()

        m += len(label_hat)
        for x in range(len(label_hat)):
            RULr = label[x][0]
            RULp = label_hat[x][0]
            pltx.append(RULr)
            plty.append(RULp)

            sum_RMSE_loss += (RULr - RULp) ** 2
            sum_MAE_loss += abs(RULr - RULp)
            E = (RULr - RULp) * 100 / RULr
            if E <= 0:
                sum_RUL_score += math.exp(-temp_num * (E / 5))
            else:
                sum_RUL_score += math.exp(temp_num * (E / 20))
            # 保存数据
    sum_RMSE_loss = math.sqrt(sum_RMSE_loss / m)
    sum_MAE_loss = sum_MAE_loss / m
    sum_RUL_score /= m
    print('RMSE:{}, MAE:{}, RUL_SCORE:{}'.format(sum_RMSE_loss, sum_MAE_loss, sum_RUL_score))
    # 保存图片
    plt.clf()
    plt.plot(range(len(pltx)), pltx, c='blue', label='RUL truth')
    plt.scatter(range(len(pltx)), plty, c='red', s=0.2, label='RUL predict')
    plt.title('Bearing1_' + str(test_bearing))
    plt.legend()
    plt.show()
    plt.clf()
